import socket, hashlib, base64, time
from traceback import print_exc

class Header:
    def __init__(self, method, path, http_version) -> None:
        super().__init__()
        self.method = method
        self.path = path
        self.http_version = http_version

class WSHandler:
    def __init__(self, sock: socket.socket) -> None:
        self.sock = sock
        self.hand_shake()
    
    def read_pack(self, pack: str):
        lines = pack.splitlines()
        header = lines[0].split(' ')
        # header = Header(b[0], b[1], b[2])
        data = {}
        for i in lines[1:]:
            a = i.split(':')
            data[a[0]] = ':'.join(a[1:]).strip()
        return header, data
    
    def make_pack(self, header:list, data:dict):
        r = ''
        r+= ' '.join(header)+'\r\n'
        for k,v in data.items():
            r+= f'{k}: {v}\r\n'
        return r

    def hand_shake(self):
        req_header, req_data = self.read_pack(self.sock.recv(1024).decode('utf-8'))
        if req_header[0] != 'GET':return
        try:
            if req_data['Connection'] != 'Upgrade':return
            if req_data['Upgrade'] != 'websocket':return
        except KeyError:return
        if 'Sec-WebSocket-Key' not in req_data.keys():return

        key = base64.b64encode(hashlib.sha1((req_data['Sec-WebSocket-Key'] + '258EAFA5-E914-47DA-95CA-C5AB0DC85B11').encode('utf-8')).digest()).decode('utf-8')

        resp_header = ['HTTP/1.1', '101', 'Switching Protocols']
        resp_data = {
            "Upgrade": "websocket",
            "Connection": "Upgrade",
            "Sec-WebSocket-Key": key
        }
        self.sock.sendall(self.make_pack(resp_header, resp_data).encode('utf-8'))
        self.sock.sendall(b"hello, world!")
        # print(self.sock.recv(1024).decode('utf-8'))

server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
while True:
    try:
        server.bind(("0.0.0.0", 43999))
        break
    except OSError:
        time.sleep(10)
print("binded!")
server.listen(10)
sock, addr = server.accept()
try:
    WSHandler(sock).hand_shake()
except:
    print_exc()
    sock.close()